import random
import torch
import os
import scipy.io as io
from PIL import Image,ImageFilter,ImageDraw
import numpy as np
import h5py
from PIL import ImageStat
import cv2

def load_data(img_path,train = True):
    gt_path = img_path.replace('.jpg','.h5').replace('images','annotation')
    img = Image.open(img_path).convert('RGB')
    mask = io.loadmat("/work2/lizirui/PCL/data_annotation/annotation/train/scene_1/scene_1_mask.mat")["mask"]
    mask_torch = torch.unsqueeze(torch.from_numpy(mask),2)    
    mask_numpy = torch.cat((mask_torch,mask_torch,mask_torch),2).numpy()   
    img = img * mask_numpy        
    gt_file = h5py.File(gt_path)
    target = np.asarray(gt_file['density'])
#    target = cv2.resize(target,(target.shape[1]/2,target.shape[0]/2),interpolation = cv2.INTER_CUBIC)*4
#    img = cv2.resize(img,(img.shape[1]/2,img.shape[0]/2),interpolation = cv2.INTER_CUBIC)*4

    
    if False:
        crop_size = (img.size[0]/2,img.size[1]/2)
        if random.randint(0,9)<= -1:
          
            dx = int(random.randint(0,1)*img.size[0]*1./2)
            dy = int(random.randint(0,1)*img.size[1]*1./2)
        else:
            dx = int(random.random()*img.size[0]*1./2)
            dy = int(random.random()*img.size[1]*1./2)
                       
        img = img.crop((dx,dy,crop_size[0]+dx,crop_size[1]+dy))
        target = target[dy:crop_size[1]+dy,dx:crop_size[0]+dx]
                               
        if random.random()>0.8:
            target = np.fliplr(target)
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
               
    target = cv2.resize(target,(target.shape[1]/8,target.shape[0]/8),interpolation = cv2.INTER_CUBIC)*64
        
    return img,target
